module thin_cloud_test_mod

  !Kent et al 2014, Dynamical core model intercomparison project: Tracer transport test cases.
  ! DCMIP1-3

	use const_mod
	use namelist_mod
	use sphere_geometry_mod
	use mesh_mod
	use vert_coord_mod
	use state_mod
	use tracer_mod
	use parallel_mod

	implicit none

	private
	public thin_cloud_test_init
	public thin_cloud_test_set_ic
	public thin_cloud_test_set_wind

	real(r8), parameter :: period = 12 * 86400 
	real(r8), parameter :: u0     = pi2 * radius / period
	real(r8), parameter :: alpha  = pi / 6 
	real(r8), parameter :: lonm   = 1.5 * pi
	real(r8), parameter :: latm   = 0
	real(r8), parameter :: h0     = 2000._r8
	real(r8), parameter :: Rm     = 0.75 * pi
	real(r8), parameter :: xim    = pi / 16
	real(r8), parameter :: lonp   = pi05
	real(r8), parameter :: latp   = 0
	real(r8), parameter :: zp1    = 3050
	real(r8), parameter :: zp2    = 5050
	real(r8), parameter :: zp3    = 8200
	real(r8), parameter :: dzp1   = 1000
	real(r8), parameter :: dzp2   = 1000
	real(r8), parameter :: dzp3   = 400
	real(r8), parameter :: Rp     = 0.25 * pi

	real(r8) T0, p0

contains
	
	subroutine thin_cloud_test_init()
		
		T0 = 300._r8
		p0 = 100000._r8

	end subroutine thin_cloud_test_init

	subroutine thin_cloud_test_set_ic(mesh, tracer)

		type(mesh_type), intent(in) :: mesh
		type(tracer_type), intent(inout) :: tracer
		integer i, j, k, itracer
		real(r8) lon, lat, p, z, r_m, zs, rz1, rz2, rz3, r_p
		real(r8),allocatable :: ps(:,:)
    real(r8),allocatable :: q1(:,:,:), q2(:,:,:), q3(:,:,:)

    allocate(ps(mesh%full_lon_ibeg:mesh%full_lon_iend, &
								mesh%full_lat_ibeg:mesh%full_lat_iend))
    allocate(q1(mesh%full_lon_ibeg:mesh%full_lon_iend, &
    						mesh%full_lat_ibeg:mesh%full_lat_iend, &
    						mesh%full_lev_ibeg:mesh%full_lev_iend))
    allocate(q2(mesh%full_lon_ibeg:mesh%full_lon_iend, &
    						mesh%full_lat_ibeg:mesh%full_lat_iend, &
    						mesh%full_lev_ibeg:mesh%full_lev_iend))
    allocate(q3(mesh%full_lon_ibeg:mesh%full_lon_iend, &
    						mesh%full_lat_ibeg:mesh%full_lat_iend, &
    						mesh%full_lev_ibeg:mesh%full_lev_iend)) 

		if (ntracers /= 1) then
			print*, 'ntracers must be 1!'
			stop
		end if

		do j = mesh%full_lat_ibeg, mesh%full_lat_iend
			lat = mesh%full_lat(j)
			do i = mesh%full_lon_ibeg, mesh%full_lon_iend
				lon = mesh%full_lon(i)
				r_m = calc_distance(lon, lat, lonm, latm) / radius
				if (r_m < Rm) then
					zs = 0.5 * h0 * (1 + cos(pi * r_m / Rm)) * cos(pi * r_m / xim)**2
				else
					zs = 0
				end if
				ps(i,j) = p0 * exp(-g * zs / Rd / T0)
			end do
		end do

    do k = mesh%full_lev_ibeg, mesh%full_lev_iend
    	do j = mesh%full_lat_ibeg, mesh%full_lat_iend
    		lat = mesh%full_lat(j)
    		do i = mesh%full_lon_ibeg, mesh%full_lon_iend
    			lon = mesh%full_lon(i)
    			p = hybrid_coord_calc_ph(k, ps(i,j))
    			z = Rd * T0 / g * log(p0 / p)
    			rz1 = abs(z - zp1)
    			rz2 = abs(z - zp2)
    			rz3 = abs(z - zp3)
    			r_p = calc_distance(lon, lat, lonp, latp) / radius
    			if (rz1 < 0.5 * dzp1 .and. r_p < Rp) then
    				q1(i,j,k) = 0.25 * (1 + cos(pi2 * rz1 / dzp1)) * (1 + cos(pi * r_p / Rp))
    			else 
    				q1(i,j,k) = 0
    			end if

    			if (rz2 < 0.5 * dzp2 .and. r_p < Rp) then
    				q2(i,j,k) = 0.25 * (1 + cos(pi2 * rz2 / dzp2)) * (1 + cos(pi * r_p / Rp))
    			else 
    			  q2(i,j,k) = 0
    			end if 

					if (rz3 < 0.5 * dzp3 .and. r_p < Rp) then
						q3(i,j,k) = 1
					else
						q3(i,j,k) = 0
					end if

					tracer%q(i,j,k,1) = q1(i,j,k) + q2(i,j,k) + q3(i,j,k)			
    		end do
    	end do
    end do

    do itracer = 1, size(tracer%q, 4)
    	call fill_zonal_halo_cell(tracer%q(:,:,:,itracer), all_halo=.true.)
    	call fill_merid_halo_cell(tracer%q(:,:,:,itracer), all_halo=.true.)
    end do

    deallocate(ps, q1, q2, q3)

	end subroutine thin_cloud_test_set_ic

	subroutine thin_cloud_test_set_wind(mesh, state, time_in_seconds)

		type(mesh_type), intent(in) :: mesh
		type(state_type), intent(inout) :: state
		real(8), intent(in) :: time_in_seconds
		real(r8) lat, lon, ulon, ulat, r_m, drmdlon, drmdlat, dzs, dzsdlon, dzsdlat
		integer i, j, k

		associate (u => state%u, v => state%v, we => state%we, zs => state%zs, phs => state%phs)

		do j = mesh%full_lat_ibeg, mesh%full_lat_iend
			lat = mesh%full_lat(j)
			do i = mesh%full_lon_ibeg, mesh%full_lon_iend
				lon = mesh%full_lon(i)
				r_m = calc_distance(lon, lat, lonm, latm) / radius
				if (r_m < Rm) then
					zs(i,j) = 0.5 * h0 * (1 + cos(pi * r_m / Rm)) * cos(pi * r_m / xim)**2
				else
					zs(i,j) = 0
				end if
				phs(i,j) = p0 * exp(-g * zs(i,j) / Rd / T0)
			end do
		end do

		do k = mesh%full_lev_ibeg, mesh%full_lev_iend
			do j = mesh%full_lat_ibeg+1, mesh%full_lat_iend-1
				lat = mesh%full_lat(j)
				do i = mesh%half_lon_ibeg, mesh%half_lon_iend
					lon = mesh%half_lon(i)
					u(i,j,k) = u0 * (cos(lat) * cos(alpha) + sin(lat) * cos(lon) * sin(alpha))
				end do
			end do
		end do
	  call fill_zonal_halo_lon(u, all_halo=.true.)

		do k = mesh%full_lev_ibeg, mesh%full_lev_iend
			do j = mesh%half_lat_ibeg, mesh%half_lat_iend
				do i = mesh%full_lon_ibeg, mesh%full_lon_iend
					lon = mesh%full_lon(i)
					v(i,j,k) = -u0 * sin(lon) * sin(alpha)
				end do
			end do
		end do

		do k = mesh%half_lev_ibeg + 1, mesh%half_lev_iend - 1
			do j = mesh%full_lat_ibeg, mesh%full_lat_iend
				lat = mesh%full_lat(j)
				do i = mesh%full_lon_ibeg, mesh%full_lon_iend
					lon = mesh%full_lon(i)
					ulon = u0 * (cos(lat) * cos(alpha) + sin(lat) * cos(lon) * sin(alpha))
					ulat = -u0 * sin(lon) * sin(alpha)
					r_m = calc_distance(lon, lat, lonm, latm) / radius
					if (r_m < Rm) then
						if (abs(r_m) < eps .or. abs(r_m - pi) < eps .or. abs(r_m + pi) < eps) then
							drmdlon = 0
							drmdlat = 0
						else
						  drmdlon = cos(latm) * cos(lat) * sin(lon - lonm) / sqrt(1 - cos(r_m)**2)
						  drmdlat = (-sin(latm) * cos(lat) + cos(latm) * sin(lat) * cos(lon - lonm)) / sqrt(1 - cos(r_m)**2)
						end if
						dzs = (-h0 * pi * 0.5 / Rm * sin(pi * r_m / Rm) * cos(pi * r_m / xim)**2 - &
						        h0 * pi / xim * (1 + cos(pi * r_m / Rm)) * cos(pi * r_m / xim) * sin(pi * r_m / xim))
						dzsdlon = dzs * drmdlon					 
						dzsdlat = dzs * drmdlat							
					else
						dzsdlon = 0
						dzsdlat = 0
					end if
					we(i,j,k) = (ulon / radius / cos(lat) * dzsdlon + ulat / radius * dzsdlat) * &
					            hybi(k) * p0 * exp(-g * zs(i,j) / Rd / T0) * (g / Rd / T0)
				end do
			end do
		end do
		end associate
	end subroutine thin_cloud_test_set_wind

end module thin_cloud_test_mod